import numpy as np 
import pandas as pd
import os
import matplotlib.pyplot as plt
import mplfinance

class Ret_calc():
    def __init__(self, args, buy_pos, sell_pos, hold_pos,  money_df, name) -> None:
        self.args = args
        self.root = os.path.join(args.result_root, name, f'loss{args.stop_loss}_ret{args.stop_ret}_open_thred{args.open_thred}_close_thred{args.close_thred}_holdnum{args.hold_num}')
        if not os.path.exists(self.root):
            os.makedirs(self.root)
        self.codes = buy_pos.columns
        # self.open = open
        # self.close = close 
        self.money_df = money_df
        self.sell_pos = sell_pos
        self.buy_pos = buy_pos 
        self.hold_pos = hold_pos
        self.get_price()
        self.lagret = self.close / self.close.shift(1) - 1
        self.cost_ratio = args.cost_ratio
        self.info = pd.Series()

    def calc_ret(self): # 添加手持金额    
        num_s = self.hold_pos.sum(axis=0)
        num_s = num_s.sort_values(ascending=False)
        plot_code = num_s.index[:10]
        for i in plot_code:
            self.plot_pos(i)
        cost_df = abs(self.hold_pos.diff(1)) * self.cost_ratio
        cost_df.fillna(0, inplace = True)
        asset = (self.close * self.hold_pos*100 - cost_df).sum(axis=1) + self.money_df.iloc[:,0]
        cumret = asset/self.args.money_init - 1  
        self.plot_cumret(cumret)
        self.turn_calc()
        self.win_lose_calc()
        self.info.to_csv(f'{self.root}/info.csv')
        
    def get_price(self):
        def read_data(root):
            data = pd.read_pickle(root).astype(float)
            data.index = pd.to_datetime(data.index)
            data = data.loc[self.hold_pos.index]
            data = data.fillna(method='ffill')
            return data
        self.open = read_data(r'data/stock_data/daily/consentrate_price/open.pkl.gzip') 
        self.close = read_data(r'data/stock_data/daily/consentrate_price/close.pkl.gzip')
        self.high = read_data(r'data/stock_data/daily/consentrate_price/high.pkl.gzip')
        self.low = read_data(r'data/stock_data/daily/consentrate_price/low.pkl.gzip')
        
    def plot_pos(self, code):
        pos_up = np.sign(self.buy_pos) * self.open
        pos_sub = np.sign(self.sell_pos) * self.close
        df = pd.DataFrame()
        df['pos+'] = pos_up[code]
        df['pos-'] = pos_sub[code]
        df['close'] = self.close[code]
        df['open'] = self.open[code]
        # df = df.iloc[].copy()
        df.replace([0], np.nan, inplace=True)

        # try:
        #     bar_values = np.array([self.open[[code]].values.squeeze(-1), self.close[[code]].values.squeeze(-1), self.high[[code]].values.squeeze(-1), self.low[[code]].values.squeeze(-1)]).T
        #     adp =[ mplfinance.make_addplot(df[['pos+']], scatter=True, color='orange'), mplfinance.make_addplot(df[['pos-']], scatter=True, color='b') ]
        #     my_color = mplfinance.make_marketcolors(up='r',
        #                             down='g',
        #                             edge='inherit',
        #                             wick='inherit',
        #                             volume='inherit')
        #     # 设置图表的背景色
        #     my_style = mplfinance.make_mpf_style(marketcolors=my_color,
        #                         figcolor='(0.82, 0.83, 0.85)',
        #                         gridcolor='(0.82, 0.83, 0.85)')

        #     price = pd.DataFrame(bar_values,index=df.index, columns=['open','close','high','low'])
        #     mplfinance.plot(data=price, type='candle', style=my_style, addplot =adp)
        #     plt.xticks(rotation=30)
        #     plt.savefig(f'{self.root}/pos_{code}_candle.png')
        #     plt.close()

        # except Exception:
        #     print('candel plot faild')

        df[['close', 'pos+', 'pos-']].plot(figsize=(16,8),style=[ '-','o','D'],alpha=0.5, color=['orange','r','g'])
        for a, b in zip(df.index, df['open'].values):
            buy_num = self.buy_pos[[code]].loc[a].values[0]   
            if buy_num != 0:
                plt.text(a, b, buy_num, ha='center', va='bottom', fontsize=10)
        for a, b in zip(df.index, df['close'].values):        
            sell_num = self.sell_pos[[code]].loc[a].values[0]
            if sell_num != 0:
                plt.text(a, b, sell_num, ha='center', va='bottom', fontsize=10)
        plt.savefig(f'{self.root}/pos_{code}_pos.png')
        plt.close()
        

    def plot_cumret(self, cumret):
        ret_all = cumret - cumret.shift(1)
        ret_all.fillna(0, inplace=True)
        df = pd.DataFrame()
        df['cumret'] = cumret
        df['back'] = df['cumret'] - df['cumret'].expanding(min_periods=1).max()
        df.plot(figsize=(16,16), subplots=True)
        plt.savefig(f'{self.root}/cumret.png')
        plt.close()
        try:
            sp = ret_all.mean()/ret_all.std() * np.sqrt(252)
        except Exception:
            sp=0
        print(f'的夏普为{sp}')
        back = df['back'].min()
        print(f'最大回撤：{back}')
        self.info['sp'] = sp
        self.info['back'] = back

    def turn_calc(self):
        X_diff = self.hold_pos.diff(1)
        X_diff_mean = abs(X_diff).mean().mean()
        X_mean = np.abs(self.hold_pos).mean().mean()
        turnover = np.divide(2 * X_mean , X_diff_mean)
        print('换手周期', turnover)
        self.info['turnover'] = turnover

    def win_lose_calc(self):
        zero_ary = np.zeros(self.sell_pos.values.shape)
        sell_cost_df = pd.DataFrame(zero_ary, index=self.sell_pos.index, columns = self.sell_pos.columns)
        win_num_df = sell_cost_df.copy()      
        for idx, idv in enumerate(self.sell_pos.index):
            for code in self.sell_pos.columns:
                sell_num = self.sell_pos.loc[idv, code]
                if sell_num > 0:
                    idv_pre = self.sell_pos.index[idx-1]
                    hold_num = self.hold_pos.loc[idv_pre, code]
                    buy_df_i = self.buy_pos.iloc[:idx]
                    buy_df_i = (buy_df_i.iloc[::-1])[code]
                    buy_i_cum = buy_df_i.cumsum()
                    sart_id = np.where(buy_i_cum>=hold_num)[0][0]
                    start_idv = buy_i_cum.index[sart_id]
                    cost_df = self.open.loc[start_idv:idv_pre, code]* np.sign(self.buy_pos.loc[start_idv:idv_pre, code])                   
                    sell_price = self.close.loc[idv,code]
                    win_num_i = sell_price > (cost_df[cost_df>0])
                    win_num_df.loc[idv,code] = np.sum(win_num_i*self.buy_pos.loc[cost_df.index, code])
                    cost_i = (cost_df*self.buy_pos.loc[start_idv:idv_pre, code]).sum()/(self.buy_pos.loc[start_idv:idv_pre, code].sum())
                    sell_cost_df.loc[idv,code] = cost_i
        win_ratio = win_num_df.sum().sum()/self.sell_pos.sum().sum()
        win_money_v = (self.sell_pos*(self.close-sell_cost_df)).values
        win_mean = np.nansum(win_money_v[win_money_v>0]*100)/ win_num_df.sum().sum()
        lose_mean = np.nansum(win_money_v[win_money_v<0]*100)/ (self.sell_pos.sum().sum() - win_num_df.sum().sum())
        win_lose_ratio = abs(win_mean/lose_mean)
        print(f'胜率:{win_ratio}， 盈亏比:{win_lose_ratio}')
        self.info['win_ratio'] = win_ratio
        self.info['win_lose_ratio'] = win_lose_ratio





#补充计算胜率与盈亏比
# 胜率 = 盈利单子数/总单子数
# 盈亏比 = 盈利总额/盈利单子数

class Ret_calc_target(Ret_calc):
    def __init__(self, args, buy_pos, sell_pos, hold_pos, money_df, name) -> None:
        super().__init__(args, buy_pos, sell_pos, hold_pos, money_df, name)
        self.root = os.path.join(args.result_root, name, f'loss{args.stop_loss}_ret{args.stop_ret}_thred{args.open_thred}_holdnum{args.hold_num}_holdt{args.hold_target}')